#pragma once

#include <gfx/dx12/dx_common.h>
#include <utl/deque.h>
#include <mutex>

namespace primal
{
	namespace gfx
	{
		template <D3D12_DESCRIPTOR_HEAP_TYPE type,u32 capacity,bool visibility>
		class DescriptorHeap : public ComPtr<ID3D12DescriptorHeap>
		{
		public:
			DescriptorHeap(ID3D12DeviceN* device);
			void allocator(descriptor_handle& handle);
			void free(descriptor_handle handle);
		private:
			utl::queue<descriptor_handle> free_handles{};
			u32 descriptor_size;
			u32 size;
			std::mutex locker{};
		};


		template<D3D12_DESCRIPTOR_HEAP_TYPE type, u32 capacity, bool visibility>
		inline DescriptorHeap<type, capacity, visibility>::DescriptorHeap(ID3D12DeviceN* device) :ComPtr<ID3D12DescriptorHeap>{}, descriptor_size{ device->GetDescriptorHandleIncrementSize(type) }, size{}
		{
			D3D12_DESCRIPTOR_HEAP_DESC desc{};
			desc.Type = type;
			desc.NumDescriptors = capacity;
			desc.Flags = visibility ? D3D12_DESCRIPTOR_HEAP_FLAG_SHADER_VISIBLE : D3D12_DESCRIPTOR_HEAP_FLAG_NONE;
			desc.NodeMask = 0u;

			DXASSERT(device->CreateDescriptorHeap(&desc, IID_PPV_ARGS(this->GetAddressOf())));
		}

		template<D3D12_DESCRIPTOR_HEAP_TYPE type, u32 capacity, bool visibility>
		inline void DescriptorHeap<type, capacity, visibility>::allocator(descriptor_handle& handle)
		{
			std::lock_guard _l{ this->locker };
			if (this->free_handles.empty())
			{
				assert(this->size < capacity);
				handle.cpu = this->ptr_->GetCPUDescriptorHandleForHeapStart();
				if (visibility)
				{
					handle.gpu = this->ptr_->GetGPUDescriptorHandleForHeapStart();
				}
				auto offset{ this->size * this->descriptor_size };
				handle.cpu.ptr += offset;
				handle.gpu.ptr += offset;
				++this->size;
			}
			else
			{
				handle = this->free_handles.front();
				this->free_handles.pop();
			}
		}

		template<D3D12_DESCRIPTOR_HEAP_TYPE type, u32 capacity, bool visibility>
		inline void DescriptorHeap<type, capacity, visibility>::free(descriptor_handle handle)
		{
			std::lock_guard _l{ this->locker };
			this->free_handles.push(handle);
		}

		using SRVDescriptorHeap = DescriptorHeap<D3D12_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV, 1024u, true>;
		using RTVDescriptorHeap = DescriptorHeap<D3D12_DESCRIPTOR_HEAP_TYPE_RTV, 1024u, false>;
		using DSVDescriptorHeap = DescriptorHeap<D3D12_DESCRIPTOR_HEAP_TYPE_DSV, 1024u, false>;
	}
}